{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem: Write a Custom Activation Function\n",
    "\n",
    "### Problem Statement\n",
    "You are tasked with implementing a **custom activation function** in PyTorch that computes the following operation: \n",
    " \n",
    "$ \\text{activation}(x) = \\tanh(x) + x $ \n",
    "\n",
    "Once implemented, this custom activation function will be used in a simple linear regression model.\n",
    "\n",
    "### Requirements\n",
    "1. **Custom Activation Function**:\n",
    "   - Implement a class `CustomActivationModel` inheriting from `torch.nn.Module`.\n",
    "   - Define the `forward` method to compute the activation function \\( \\text{tanh}(x) + x \\).\n",
    "\n",
    "2. **Integration with Linear Regression**:\n",
    "   - Use the custom activation function in a simple linear regression model.\n",
    "   - The model should include:\n",
    "     - A single linear layer.\n",
    "     - The custom activation function applied to the output of the linear layer.\n",
    "\n",
    "### Constraints\n",
    "- The custom activation function should not have any learnable parameters.\n",
    "- Ensure compatibility with PyTorch tensors for forward pass computations.\n",
    "\n",
    "<details>\n",
    "  <summary>💡 Hint</summary>\n",
    "  Some details: https://stackoverflow.com/questions/55765234/pytorch-custom-activation-functions\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate synthetic data\n",
    "torch.manual_seed(42)\n",
    "X = torch.rand(100, 1) * 10  # 100 data points between 0 and 10\n",
    "y = 2 * X + 3 + torch.randn(100, 1)  # Linear relationship with noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'NoneType' object has no attribute 'size'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 25\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m epoch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(epochs):\n\u001b[1;32m     23\u001b[0m     \u001b[38;5;66;03m# Forward pass\u001b[39;00m\n\u001b[1;32m     24\u001b[0m     predictions \u001b[38;5;241m=\u001b[39m model(X)\n\u001b[0;32m---> 25\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[43mcriterion\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpredictions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     27\u001b[0m     \u001b[38;5;66;03m# Backward pass and optimization\u001b[39;00m\n\u001b[1;32m     28\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n",
      "File \u001b[0;32m~/Library/Caches/pypoetry/virtualenvs/mylogs--5zRa99S-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1734\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1736\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Library/Caches/pypoetry/virtualenvs/mylogs--5zRa99S-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1745\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1746\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1747\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m   1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n",
      "File \u001b[0;32m~/Library/Caches/pypoetry/virtualenvs/mylogs--5zRa99S-py3.10/lib/python3.10/site-packages/torch/nn/modules/loss.py:608\u001b[0m, in \u001b[0;36mMSELoss.forward\u001b[0;34m(self, input, target)\u001b[0m\n\u001b[1;32m    607\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor, target: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 608\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmse_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreduction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreduction\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Library/Caches/pypoetry/virtualenvs/mylogs--5zRa99S-py3.10/lib/python3.10/site-packages/torch/nn/functional.py:3781\u001b[0m, in \u001b[0;36mmse_loss\u001b[0;34m(input, target, size_average, reduce, reduction)\u001b[0m\n\u001b[1;32m   3771\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_variadic(\u001b[38;5;28minput\u001b[39m, target):\n\u001b[1;32m   3772\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m   3773\u001b[0m         mse_loss,\n\u001b[1;32m   3774\u001b[0m         (\u001b[38;5;28minput\u001b[39m, target),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   3779\u001b[0m         reduction\u001b[38;5;241m=\u001b[39mreduction,\n\u001b[1;32m   3780\u001b[0m     )\n\u001b[0;32m-> 3781\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (target\u001b[38;5;241m.\u001b[39msize() \u001b[38;5;241m==\u001b[39m \u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msize\u001b[49m()):\n\u001b[1;32m   3782\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(\n\u001b[1;32m   3783\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUsing a target size (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtarget\u001b[38;5;241m.\u001b[39msize()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m) that is different to the input size (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39msize()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m). \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   3784\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThis will likely lead to incorrect results due to broadcasting. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   3785\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPlease ensure they have the same size.\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   3786\u001b[0m         stacklevel\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m,\n\u001b[1;32m   3787\u001b[0m     )\n\u001b[1;32m   3788\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size_average \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m reduce \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'NoneType' object has no attribute 'size'"
     ]
    }
   ],
   "source": [
    "# Define the Linear Regression Model\n",
    "class CustomActivationModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CustomActivationModel, self).__init__()\n",
    "        self.linear = nn.Linear(1, 1)  # Single input and single output\n",
    "\n",
    "    # TODO: Implement the forward pass\n",
    "    def custom_activation(self, x):\n",
    "        ...\n",
    "            \n",
    "    # TODO: Implement the forward pass\n",
    "    def forward(self, x):\n",
    "        ...\n",
    "\n",
    "# Initialize the model, loss function, and optimizer\n",
    "model = CustomActivationModel()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "# Training loop\n",
    "epochs = 1000\n",
    "for epoch in range(epochs):\n",
    "    # Forward pass\n",
    "    predictions = model(X)\n",
    "    loss = criterion(predictions, y)\n",
    "\n",
    "    # Backward pass and optimization\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # Log progress every 100 epochs\n",
    "    if (epoch + 1) % 100 == 0:\n",
    "        print(f\"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learned weight: 1.9557, Learned bias: 2.2181\n",
      "Predictions for [[4.0], [7.0]]: [[11.04088020324707], [16.907970428466797]]\n"
     ]
    }
   ],
   "source": [
    "# Display the learned parameters\n",
    "[w, b] = model.linear.parameters()\n",
    "print(f\"Learned weight: {w.item():.4f}, Learned bias: {b.item():.4f}\")\n",
    "\n",
    "# Testing on new data\n",
    "X_test = torch.tensor([[4.0], [7.0]])\n",
    "with torch.no_grad():\n",
    "    predictions = model(X_test)\n",
    "    print(f\"Predictions for {X_test.tolist()}: {predictions.tolist()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
